{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Bistable System w/ Gaussian Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pydrake.all import (AddRandomInputs, BasicVector, DiagramBuilder,\n",
    "                         LeafSystem, PortDataType, PyPlotVisualizer,\n",
    "                         RandomDistribution, Simulator)\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "\n",
    "def dynamics(x, w):\n",
    "    return x - x**3 + .3 * w\n",
    "\n",
    "\n",
    "class SimpleStochasticSystem(LeafSystem):\n",
    "\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareInputPort(\"noise\", PortDataType.kVectorValued, 1,\n",
    "                              RandomDistribution.kGaussian)\n",
    "        self.DeclareContinuousState(1)\n",
    "        self.DeclareVectorOutputPort(BasicVector(1), self.CopyStateOut)\n",
    "\n",
    "    # xdot(t) = x(t) - x^3(t) + w(t)\n",
    "    def DoCalcTimeDerivatives(self, context, derivatives):\n",
    "        x = context.get_continuous_state_vector().GetAtIndex(0)\n",
    "        w = self.EvalVectorInput(context, 0).GetAtIndex(0)\n",
    "        xdot = dynamics(x, w)\n",
    "        derivatives.get_mutable_vector().SetAtIndex(0, xdot)\n",
    "\n",
    "    # y(t) = x(t)\n",
    "    def CopyStateOut(self, context, output):\n",
    "        x = context.get_continuous_state_vector().get_value()\n",
    "        y = output.get_mutable_value()\n",
    "        y[:] = x\n",
    "\n",
    "\n",
    "# Note: This is a candidate for moving to a more central location.\n",
    "class HistogramVisualizer(PyPlotVisualizer):\n",
    "\n",
    "    def __init__(self,\n",
    "                 num_samples,\n",
    "                 bins,\n",
    "                 xlim,\n",
    "                 ylim,\n",
    "                 draw_timestep,\n",
    "                 figsize=None):\n",
    "        PyPlotVisualizer.__init__(self, draw_timestep, figsize=figsize)\n",
    "        for i in range(0, num_samples):\n",
    "            self.DeclareInputPort(PortDataType.kVectorValued, 1)\n",
    "        self.num_samples = num_samples\n",
    "        self.bins = bins\n",
    "        self.data = [0]*num_samples\n",
    "        self.limits = xlim\n",
    "        self.ax.set_xlim(xlim)\n",
    "        self.ax.axis(\"auto\")\n",
    "        self.ax.set_ylim(ylim)\n",
    "        self.patches = None\n",
    "\n",
    "    def draw(self, context):\n",
    "        if (self.patches):\n",
    "            t = [p.remove() for p in self.patches]\n",
    "        for i in range(0, self.num_samples):\n",
    "            self.data[i] = self.EvalVectorInput(context, i).GetAtIndex(0)\n",
    "        count, bins, self.patches = self.ax.hist(\n",
    "            self.data, bins=self.bins, range=self.limits, density=True,\n",
    "            facecolor=\"b\")\n",
    "        self.ax.set_title(\"t = \" + str(context.get_time()))\n",
    "\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "num_particles = 1000\n",
    "num_bins = 100\n",
    "xlim = [-2, 2]\n",
    "ylim = [-1, 3.5]\n",
    "draw_timestep = .25\n",
    "visualizer = builder.AddSystem(\n",
    "    HistogramVisualizer(num_particles, num_bins, xlim, ylim, draw_timestep))\n",
    "x = np.linspace(xlim[0], xlim[1], 100)\n",
    "visualizer.ax.plot(x, dynamics(x, 0), \"k\", linewidth=2)\n",
    "\n",
    "for i in range(0, num_particles):\n",
    "    sys = builder.AddSystem(SimpleStochasticSystem())\n",
    "    builder.Connect(sys.get_output_port(0), visualizer.get_input_port(i))\n",
    "\n",
    "AddRandomInputs(.1, builder)\n",
    "\n",
    "diagram = builder.Build()\n",
    "simulator = Simulator(diagram)\n",
    "simulator.get_mutable_integrator().set_fixed_step_mode(True)\n",
    "simulator.get_mutable_integrator().set_maximum_step_size(0.1)\n",
    "\n",
    "AdvanceToAndVisualize(simulator, visualizer, 20.0, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Stochastic Van der Pol Oscillator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pydrake.all import (AddRandomInputs, BasicVector, DiagramBuilder,\n",
    "                         LeafSystem, PortDataType, PyPlotVisualizer,\n",
    "                         RandomDistribution, Simulator)\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "\n",
    "class VanDerPolParticles(LeafSystem):\n",
    "\n",
    "    def __init__(self, num_particles, mu=1.0):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareInputPort(\"noise\", PortDataType.kVectorValued,\n",
    "                              num_particles, RandomDistribution.kGaussian)\n",
    "        self.DeclareContinuousState(num_particles, num_particles, 0)\n",
    "        self.DeclareVectorOutputPort(BasicVector(2 * num_particles),\n",
    "                                     self.CopyStateOut)\n",
    "        self.num_particles = num_particles\n",
    "        self.mu = mu\n",
    "\n",
    "    # TODO(russt):  SetRandomState to  [-0.1144;2.0578] + 0.01*randn(...)\n",
    "\n",
    "    def DoCalcTimeDerivatives(self, context, derivatives):\n",
    "        # TODO(russt):  Update this to get_position/velocity once those are\n",
    "        # bound.\n",
    "        x = context.get_continuous_state_vector().CopyToVector()\n",
    "        q = x[:self.num_particles]\n",
    "        qdot = x[self.num_particles:]\n",
    "        w = self.EvalVectorInput(context, 0).CopyToVector()\n",
    "        qddot = -self.mu * (q * q - 1) * qdot - q + .5 * w\n",
    "        derivatives.get_mutable_vector().SetFromVector(\n",
    "            np.concatenate((qdot, qddot)))\n",
    "\n",
    "    # y(t) = x(t)\n",
    "    def CopyStateOut(self, context, output):\n",
    "        x = context.get_continuous_state_vector().get_value()\n",
    "        y = output.get_mutable_value()\n",
    "        y[:] = x\n",
    "\n",
    "\n",
    "# Note: This is a candidate for moving to a more central location.\n",
    "class Particle2DVisualizer(PyPlotVisualizer):\n",
    "\n",
    "    def __init__(self, num_particles, xlim, ylim, draw_timestep):\n",
    "        PyPlotVisualizer.__init__(self, draw_timestep)\n",
    "        self.DeclareInputPort(PortDataType.kVectorValued, 2 * num_particles)\n",
    "        self.num_particles = num_particles\n",
    "        self.ax.set_xlim(xlim)\n",
    "        self.ax.set_ylim(ylim)\n",
    "        zero = [0]*num_particles\n",
    "        self.lines, = self.ax.plot(zero, zero, \"b.\")\n",
    "\n",
    "    def draw(self, context):\n",
    "        xy = self.EvalVectorInput(context, 0).CopyToVector()\n",
    "        self.lines.set_xdata(xy[:self.num_particles])\n",
    "        self.lines.set_ydata(xy[self.num_particles:])\n",
    "        self.ax.set_title(\"t = \" + str(context.get_time()))\n",
    "\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "num_particles = 5000\n",
    "xlim = [-2.75, 2.75]\n",
    "ylim = [-3.25, 3.25]\n",
    "draw_timestep = .25\n",
    "sys = builder.AddSystem(VanDerPolParticles(num_particles))\n",
    "visualizer = builder.AddSystem(\n",
    "    Particle2DVisualizer(num_particles, xlim, ylim, draw_timestep))\n",
    "builder.Connect(sys.get_output_port(0), visualizer.get_input_port(0))\n",
    "AddRandomInputs(.1, builder)\n",
    "\n",
    "# TODO(russt): Plot nominal limit cycle.\n",
    "\n",
    "diagram = builder.Build()\n",
    "simulator = Simulator(diagram)\n",
    "simulator.set_publish_every_time_step(False)\n",
    "simulator.get_mutable_integrator().set_fixed_step_mode(True)\n",
    "simulator.get_mutable_integrator().set_maximum_step_size(0.1)\n",
    "\n",
    "AdvanceToAndVisualize(simulator, visualizer, 20.0, 1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
